package com.example.mcp.server.service;

import static org.springframework.ai.chat.memory.ChatMemory.CONVERSATION_ID;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import com.alibaba.cloud.ai.memory.redis.RedisChatMemoryRepository;
import com.example.mcp.server.model.McpInitializeRequest;
import com.example.mcp.server.model.McpInitializeResponse;
import com.example.mcp.server.model.McpListToolsResponse;
import com.example.mcp.server.model.McpTool;
import com.example.mcp.server.model.McpToolCallRequest;
import com.example.mcp.server.model.McpToolCallResponse;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import javax.script.ScriptEngine;
import javax.script.ScriptEngineManager;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

@Service
public class McpService {

    @Autowired
    private OpenAiChatModel openAiChatModel;
    @Autowired
    private DashScopeChatModel dashScopeChatModel;
    @Autowired
    private RedisChatMemoryRepository redisChatMemoryRepository;

    private final Map<String, McpTool> tools = new ConcurrentHashMap<>();
    private final Map<String, String> sessions = new ConcurrentHashMap<>();

    public McpService() {

        initializeTools();
    }

    private void initializeTools() {
        // Add weather tool
        tools.put("get_weather", new McpTool(
                "get_weather",
                "Get current weather information for a given location",
                Map.of(
                        "type", "object",
                        "properties", Map.of(
                                "location", Map.of(
                                        "type", "string",
                                        "description", "The city and country (e.g., 'London, UK')"
                                ),
                                "unit", Map.of(
                                        "type", "string",
                                        "enum", List.of("celsius", "fahrenheit"),
                                        "description", "Temperature unit"
                                )
                        ),
                        "required", List.of("location")
                )
        ));

        // Add calculator tool
        tools.put("calculate", new McpTool(
                "calculate",
                "Perform mathematical calculations",
                Map.of(
                        "type", "object",
                        "properties", Map.of(
                                "expression", Map.of(
                                        "type", "string",
                                        "description", "Mathematical expression to evaluate"
                                )
                        ),
                        "required", List.of("expression")
                )
        ));

        // Add AI assistant tool
        tools.put("ai_assistant", new McpTool(
                "ai_assistant",
                "Get AI assistance for general questions and tasks using DeepSeek Chat",
                Map.of(
                        "type", "object",
                        "properties", Map.of(
                                "query", Map.of(
                                        "type", "string",
                                        "description", "The question or task for the DeepSeek AI assistant"
                                ),
                                "context", Map.of(
                                        "type", "string",
                                        "description", "Additional context for the query"
                                )
                        ),
                        "required", List.of("query")
                )
        ));
    }

    public McpInitializeResponse initialize(McpInitializeRequest request) {
        String sessionId = UUID.randomUUID().toString();
        sessions.put(sessionId, "initialized");

        McpInitializeResponse response = new McpInitializeResponse();
        response.setId(request.getId());
        return response;
    }

    public McpListToolsResponse listTools(String sessionId) {
        if (!sessions.containsKey(sessionId)) {
            throw new IllegalArgumentException("Invalid session ID");
        }

        McpListToolsResponse response = new McpListToolsResponse();
        response.setTools(new ArrayList<>(tools.values()));
        return response;
    }

    public List<McpTool> getAllTools() {
        return new ArrayList<>(tools.values());
    }

    public McpToolCallResponse callTool(McpToolCallRequest request) {
        String toolName = request.getName();
        Map<String, Object> arguments = (Map<String, Object>) request.getArguments();

        if (!tools.containsKey(toolName)) {
            McpToolCallResponse errorResponse = new McpToolCallResponse();
            errorResponse.setContent(Map.of(
                    "error", "Tool not found: " + toolName,
                    "available_tools", tools.keySet()
            ));
            errorResponse.setIsError(true);
            return errorResponse;
        }

        try {
            Object result = executeTool(toolName, arguments);

            McpToolCallResponse response = new McpToolCallResponse();
            response.setContent(result);
            response.setIsError(false);
            return response;

        } catch (Exception e) {
            McpToolCallResponse errorResponse = new McpToolCallResponse();
            errorResponse.setContent(Map.of(
                    "error", "Tool execution failed: " + e.getMessage()
            ));
            errorResponse.setIsError(true);
            return errorResponse;
        }
    }

    private Object executeTool(String toolName, Map<String, Object> arguments) {
        switch (toolName) {
            case "get_weather":
                return getWeather(arguments);
            case "calculate":
                return calculate(arguments);
            case "ai_assistant":
                return aiAssistant(arguments);
            default:
                throw new IllegalArgumentException("Unknown tool: " + toolName);
        }
    }

    private Object getWeather(Map<String, Object> arguments) {
        String location = (String) arguments.get("location");
        String unit = (String) arguments.getOrDefault("unit", "celsius");

        // Mock weather data - in real implementation, call weather API
        Map<String, Object> weatherData = Map.of(
                "location", location,
                "temperature", unit.equals("celsius") ? 22.5 : 72.5,
                "unit", unit,
                "condition", "Partly cloudy",
                "humidity", 65,
                "wind_speed", "12 km/h"
        );

        return weatherData;
    }

    private Object calculate(Map<String, Object> arguments) {
        String expression = (String) arguments.get("expression");

        try {
            // Simple expression evaluation - in production use proper math parser
            double result = evaluateExpression(expression);
            return Map.of(
                    "expression", expression,
                    "result", result
            );
        } catch (Exception e) {
            throw new RuntimeException("Invalid expression: " + expression);
        }
    }

    private double evaluateExpression(String expression) {
        try {
            // Use JavaScript engine for proper expression evaluation
            ScriptEngineManager mgr = new ScriptEngineManager();
            ScriptEngine engine = mgr.getEngineByName("JavaScript");

            if (engine == null) {
                throw new IllegalStateException("JavaScript engine not available");
            }

            // Clean up the expression
            expression = expression.trim();
            if (!expression.matches("^[0-9+\\-*/().\\s]+$")) {
                throw new IllegalArgumentException("Invalid characters in expression");
            }

            Object result = engine.eval(expression);
            return Double.parseDouble(result.toString());

        } catch (Exception e) {
            // Provide better error messages
            throw new IllegalArgumentException("Could not evaluate expression: " + expression + ". Error: " + e.getMessage());
        }
    }

    private Object aiAssistant(Map<String, Object> arguments) {
        String query = (String) arguments.get("query");
        String context = (String) arguments.getOrDefault("context", "");

        String fullQuery = context.isEmpty() ? query : context + "\n\n" + query;

        try {

            ChatMemory chatMemory = MessageWindowChatMemory.builder()
                    .chatMemoryRepository(redisChatMemoryRepository)
                    .maxMessages(100)
                    .build();
            ChatClient build = ChatClient.builder(dashScopeChatModel)
//                    .defaultAdvisors(new SimpleLoggerAdvisor())
                    .defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemory).build())
                    .build();
            String response = build.prompt()
                    .advisors(a -> a.param(CONVERSATION_ID, "111"))
                    .user(fullQuery)
                    .call()
                    .content();

            return Map.of(
                    "query", query,
                    "response", response,
                    "context", context
            );
        } catch (Exception e) {
            // Handle DeepSeek specific issues
            String errorMessage = "DeepSeek AI assistant failed: " + e.getMessage();

            // Provide more helpful error for common DeepSeek issues
            if (e.getMessage() != null) {
                if (e.getMessage().contains("401")) {
                    errorMessage = "Invalid DeepSeek API key. Please check your API key configuration.";
                } else if (e.getMessage().contains("429")) {
                    errorMessage = "DeepSeek API rate limit exceeded. Please try again later.";
                } else if (e.getMessage().contains("500") || e.getMessage().contains("503")) {
                    errorMessage = "DeepSeek service temporarily unavailable. Please try again later.";
                }
            }

            return Map.of(
                    "query", query,
                    "response", errorMessage,
                    "context", context,
                    "error", true
            );
        }
    }

    public boolean validateSession(String sessionId) {
        return sessions.containsKey(sessionId);
    }

    public String createSession() {
        String sessionId = UUID.randomUUID().toString();
        sessions.put(sessionId, "active");
        return sessionId;
    }

    public void closeSession(String sessionId) {
        sessions.remove(sessionId);
    }


}